import tushare as ts
import numpy as np
import pandas as pd
import talib
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
from sklearn.model_selection import GridSearchCV
import matplotlib.pyplot as plt

# 初始化pro接口
pro = ts.pro_api('47b4cad5b8693e659ce45d1625d20211311ccb4039d69ba828736eb4')

# 拉取数据
df = pro.daily(**{
    "ts_code": "000001.SZ",
    "trade_date": "",
    "start_date": 20200101,
    "end_date": 20241231,
    "offset": "",
    "limit": ""
}, fields=[
    "ts_code",
    "trade_date",
    "open",
    "high",
    "low",
    "close",
    "pre_close",
    "change",
    "pct_chg",
    "vol",
    "amount"
])
# 将trade_date设置为索引
df = df.set_index('trade_date')
# 生成简单衍生变量
df['close-open'] = (df['close'] - df['open']) / df['open']
df['high-low'] = (df['high'] - df['low']) / df['low']
df['price_change'] = df['close'] - df['pre_close']
df['p_change'] = df['pct_chg']

# 生成移动平均线指标MA值
df['MA5'] = df['close'].rolling(5).mean()
df['MA10'] = df['close'].rolling(10).mean()
df.dropna(inplace=True)

# 用TA-Lib库生成相对强弱指标RSI值
df['RSI'] = talib.RSI(df['close'], timeperiod=12)
# 用TA-Lib库生成动量指标MOM值
df['MOM'] = talib.MOM(df['close'], timeperiod=5)
# 用TA-Lib库生成指数移动平均值EMA
df['EMA12'] = talib.EMA(df['close'], timeperiod=12)
df['EMA26'] = talib.EMA(df['close'], timeperiod=26)
# 用TA-Lib库生成异同移动平均线MACD值
df['MACD'], df['MACDsignal'], df['MACDhist'] = talib.MACD(df['close'], fastperiod=6, slowperiod=12, signalperiod=9)
df.dropna(inplace=True)

# 提取特征变量和目标变量
X = df[['close', 'vol', 'close-open', 'MA5', 'MA10', 'high-low', 'RSI', 'MOM', 'EMA12', 'MACD', 'MACDsignal', 'MACDhist']]
y = np.where(df['price_change'].shift(-1) > 0, 1, -1)

# 划分训练集和测试集
X_length = X.shape[0]
split = int(X_length * 0.9)
X_train = X[:split].copy()
X_test = X[split:].copy()
y_train = y[:split]
y_test = y[split:]

# 模型搭建
model = RandomForestClassifier(max_depth=3, n_estimators=10, min_samples_leaf=10, random_state=123)
model.fit(X_train, y_train)

# 预测下一天的股价涨跌情况
y_pred = model.predict(X_test)
pred_df = pd.DataFrame()
pred_df['预测值'] = list(y_pred)
pred_df['实际值'] = list(y_test)
print(pred_df.head())

# 预测属于各个分类的概率
y_pred_proba = model.predict_proba(X_test)
proba_df = pd.DataFrame(y_pred_proba, columns=['分类为-1的概率', '分类为1的概率'])
print(proba_df.head())

# 模型准确度评估
accuracy = accuracy_score(y_pred, y_test)
print(f"模型准确度: {accuracy}")
print(f"模型得分: {model.score(X_test, y_test)}")

# 分析特征变量的特征重要性
importances = model.feature_importances_
importance_df = pd.DataFrame()
importance_df['特征'] = X.columns
importance_df['特征重要性'] = importances
importance_df = importance_df.sort_values('特征重要性', ascending=False)
print(importance_df)

# 参数调优
parameters = {'n_estimators': [5, 10, 20],'max_depth': [2, 3, 4, 5, 6],'min_samples_leaf': [5, 10, 20, 30]}
new_model = RandomForestClassifier(random_state=123)
grid_search = GridSearchCV(new_model, parameters, cv=6, scoring='accuracy')
grid_search.fit(X_train, y_train)
print(f"最优参数: {grid_search.best_params_}")

# 收益回测曲线绘制
X_test['prediction'] = model.predict(X_test)
X_test['p_change'] = (X_test['close'] - X_test['close'].shift(1)) / X_test['close'].shift(1)
X_test['origin'] = (X_test['p_change'] + 1).cumprod()
X_test['strategy'] = (X_test['prediction'].shift(1) * X_test['p_change'] + 1).cumprod()
X_test[['strategy', 'origin']].dropna().plot()
plt.gcf().autofmt_xdate()
plt.show()